import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import pointbiserialr, chi2_contingency
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import LabelBinarizer
from FeatureEngineering import feature_engineering, EmployeeAttritionFeatureEngineer
from src.utils.paths import RAW_DATA_DIR, PROCESSED_DATA_DIR, FIG_DIR


def data_analysis():
    # 永久设置显示选项（直到你更改它为止）
    pd.set_option('display.max_rows', None)
    pd.set_option('display.max_columns', None)
    pd.set_option('display.width', None)
    pd.set_option('display.max_colwidth', None)
    pd.set_option('display.expand_frame_repr', False)

    data = pd.read_csv(RAW_DATA_DIR / 'train.csv')
    data = pd.get_dummies(data)
    df, feature_columns = feature_engineering(data)
    engineer = EmployeeAttritionFeatureEngineer()
    df = engineer.transform(df)

    all_cols = [
        'Attrition', 'Age', 'DistanceFromHome', 'Education', 'EmployeeNumber',
        'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel',
        'JobSatisfaction', 'MonthlyIncome', 'NumCompaniesWorked',
        'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
        'StockOptionLevel', 'TotalWorkingYears',
        'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsAtCompany',
        'YearsInCurrentRole', 'YearsSinceLastPromotion', 'YearsWithCurrManager',
        'BusinessTravel_Non-Travel', 'BusinessTravel_Travel_Frequently',
        'BusinessTravel_Travel_Rarely', 'Department_Human Resources',
        'Department_Research & Development', 'Department_Sales',
        'EducationField_Human Resources', 'EducationField_Life Sciences',
        'EducationField_Marketing', 'EducationField_Medical',
        'EducationField_Other', 'EducationField_Technical Degree',
        'Gender_Female',
        'MaritalStatus_Divorced', 'MaritalStatus_Married',
        'MaritalStatus_Single', 'OverTime_Yes',
        'isManager', '开始工作年龄', '入职收入增长率', '换工作频率', '在岗时间比例', '工作稳定指数', '搭档稳定指数',
        '晋升停滞率', '整体满意度指数', '满意度差值', '绝对收入增长', '平均在职时长', '加班_满意度差值',
        '管理岗_绝对增长',
        '资历_满意度', '加班_工作生活平衡', '晋升停滞_log', '收入增长率_sqrt', '工作稳定指数_exp'
    ]

    label_col = 'Attrition'

    binary_cols = [
        'BusinessTravel_Non-Travel', 'BusinessTravel_Travel_Frequently',
        'BusinessTravel_Travel_Rarely', 'Department_Human Resources',
        'Department_Research & Development', 'Department_Sales',
        'EducationField_Human Resources', 'EducationField_Life Sciences',
        'EducationField_Marketing', 'EducationField_Medical',
        'EducationField_Other', 'EducationField_Technical Degree',
        'Gender_Female', 'MaritalStatus_Divorced', 'MaritalStatus_Married',
        'MaritalStatus_Single', 'OverTime_Yes', 'isManager'
    ]

    multi_cols = [
        'Education', 'EnvironmentSatisfaction', 'JobInvolvement', 'JobLevel',
        'JobSatisfaction', 'PerformanceRating', 'RelationshipSatisfaction',
        'StockOptionLevel', 'WorkLifeBalance'
    ]

    continuous_cols = list(set(all_cols) - set([label_col]) - set(binary_cols) - set(multi_cols))

    # print(continuous_cols)
    # 存储结果
    continuous_results = []
    binary_results = []
    multi_results = []

    ### 连续型变量分析
    for col in continuous_cols:
        try:
            r, p = pointbiserialr(df[col], df[label_col])
            auc = roc_auc_score(df[label_col], df[col])
            continuous_results.append(
                {'Feature': col, 'Type': 'Continuous', 'PointBiserial_r': r, 'p_value': p, 'AUC': auc})
        except Exception as e:
            continuous_results.append({'Feature': col, 'Type': 'Continuous', 'Error': str(e)})

    ### 二分类变量分析
    for col in binary_cols:
        try:
            r = np.corrcoef(df[col], df[label_col])[0, 1]
            auc = roc_auc_score(df[label_col], df[col])
            binary_results.append({'Feature': col, 'Type': 'Binary', 'Correlation': r, 'AUC': auc})
        except Exception as e:
            binary_results.append({'Feature': col, 'Type': 'Binary', 'Error': str(e)})

    ### 多分类变量分析
    for col in multi_cols:
        try:
            contingency = pd.crosstab(df[col], df[label_col])
            chi2, p, _, _ = chi2_contingency(contingency)

            # One-vs-All AUC（多分类特征做one-hot后算平均AUC）
            lb = LabelBinarizer()
            lb.fit(df[col])
            transformed = lb.transform(df[col])
            aucs = []
            for i in range(transformed.shape[1]):
                try:
                    auc = roc_auc_score(df[label_col], transformed[:, i])
                    aucs.append(auc)
                except:
                    continue
            mean_auc = np.mean(aucs) if aucs else np.nan

            multi_results.append({'Feature': col, 'Type': 'Multiclass', 'Chi2_p_value': p, 'Mean_AUC': mean_auc})
        except Exception as e:
            multi_results.append({'Feature': col, 'Type': 'Multiclass', 'Error': str(e)})

    # 汇总为 DataFrame
    result_df1 = pd.DataFrame(continuous_results)
    result_df2 = pd.DataFrame(binary_results)
    result_df3 = pd.DataFrame(multi_results)

    result_df1.sort_values('AUC', inplace=True)
    result_df2.sort_values('AUC', inplace=True)
    result_df3.sort_values('Chi2_p_value', inplace=True)

    print(result_df1)
    print('-' * 80)
    print(result_df2)
    print('-' * 80)
    print(result_df3)
    return result_df1, result_df2, result_df3


def plot_feature_label_correlation(continuous_results, binary_results, multi_results):
    """
    分别绘制连续变量、二分类变量、多分类变量与标签的相关性图。

    参数:
    - continuous_results: DataFrame，包含列 ['Feature', 'PointBiserial_r', 'AUC']
    - binary_results:     DataFrame，包含列 ['Feature', 'Correlation', 'AUC']
    - multi_results:      DataFrame，包含列 ['Feature', 'Chi2_p_value', 'Mean_AUC']
    """

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文可改为 SimHei
    plt.rcParams['axes.unicode_minus'] = False
    # sns.set(style='whitegrid')

    # 1. 连续型变量
    cont_df = continuous_results.dropna(subset=['PointBiserial_r', 'AUC']).copy()
    cont_df = cont_df.sort_values('PointBiserial_r', key=abs, ascending=False)
    highlight_names = ['Age', '换工作频率', 'TotalWorkingYears', 'YearsWithCurrManager', '入职收入增长率',
                       '绝对收入增长', '平均在职时长', 'YearsInCurrentRole', 'YearsAtCompany', '整体满意度指数',
                       'MonthlyIncome', '加班_工作生活平衡', '收入增长率_sqrt', '资历_满意度', ]
    colors = ['red' if feat in highlight_names else 'gray' for feat in cont_df['Feature']]

    plt.figure(figsize=(12, 9))
    sns.barplot(x='Feature', y='PointBiserial_r',
                data=cont_df, palette=colors)
    plt.axhline(0, color='gray', linestyle='--')
    plt.xticks(rotation=90)
    plt.title('连续变量与离职情况的点二列相关系数')
    plt.ylabel('Point-Biserial r')
    plt.xlabel('Feature')
    plt.tight_layout()
    plt.savefig(FIG_DIR / "连续变量与离职情况的点二列相关系数.png")
    plt.show()

    plt.figure(figsize=(12, 9))
    cont_df['AUC_Centered'] = cont_df['AUC'] - 0.5
    cont_df = cont_df.sort_values('AUC_Centered', ascending=False)
    colors = ['red' if feat in highlight_names else 'gray' for feat in cont_df['Feature']]
    sns.barplot(x='Feature', y='AUC_Centered',
                data=cont_df, palette=colors)
    plt.axhline(0, color='gray', linestyle='--')
    plt.xticks(rotation=90)
    plt.title('连续变量预测离职情况的AUC（以0.5为基线）')
    plt.ylabel('AUC - 0.5')
    plt.xlabel('Feature')
    plt.tight_layout()
    plt.savefig(FIG_DIR / "连续变量预测离职情况的AUC.png")
    plt.show()

    # 2. 二分类变量
    bin_df = binary_results.dropna(subset=['Correlation', 'AUC']).copy()
    bin_df = bin_df.sort_values('Correlation', key=abs, ascending=False)
    highlight_names = ['MaritalStatus_Single', 'OverTime_Yes', 'isManager']
    colors = ['red' if feat in highlight_names else 'gray' for feat in bin_df['Feature']]

    plt.figure(figsize=(12, 9))
    sns.barplot(x='Feature', y='Correlation',
                data=bin_df, palette=colors)
    plt.axhline(0, color='gray', linestyle='--')
    plt.xticks(rotation=90)
    plt.title('二分类变量与离职情况的皮尔逊相关系数')
    plt.ylabel('Correlation')
    plt.xlabel('Feature')
    plt.tight_layout()
    plt.savefig(FIG_DIR / "二分类变量与离职情况的皮尔逊相关系数.png")
    plt.show()

    plt.figure(figsize=(12, 9))
    bin_df['AUC_Centered'] = bin_df['AUC'] - 0.5
    bin_df = bin_df.sort_values('AUC_Centered', ascending=False)
    colors = ['red' if feat in highlight_names else 'gray' for feat in bin_df['Feature']]
    sns.barplot(x='Feature', y='AUC_Centered',
                data=bin_df, palette=colors)
    plt.axhline(0, color='gray', linestyle='--')
    plt.xticks(rotation=90)
    plt.title('二分类变量预测离职情况的AUC（以0.5为基线）')
    plt.ylabel('AUC - 0.5')
    plt.xlabel('Feature')
    plt.tight_layout()
    plt.savefig(FIG_DIR / "二分类变量预测离职情况的AUC.png")
    plt.show()

    # 3. 多分类变量
    multi_df = multi_results.dropna(subset=['Chi2_p_value', 'Mean_AUC']).copy()
    multi_df['-log10(p)'] = -np.log10(multi_df['Chi2_p_value'])
    multi_df = multi_df.sort_values('-log10(p)', ascending=False)
    highlight_names = ['StockOptionLevel', 'JobLevel']
    colors = ['red' if feat in highlight_names else 'gray' for feat in multi_df['Feature']]

    plt.figure(figsize=(12, 9))
    sns.barplot(x='Feature', y='-log10(p)',
                data=multi_df, palette=colors)
    # plt.axhline(-np.log10(0.05), color='red', linestyle='--', label='p=0.05')
    plt.xticks(rotation=90)
    plt.title('多分类变量与离职情况的显著性（-log10(p)）')
    plt.ylabel('-log10(p-value)')
    plt.xlabel('Feature')
    plt.legend()
    plt.tight_layout()
    plt.savefig(FIG_DIR / "多分类变量与离职情况的显著性.png")
    plt.show()


if __name__ == '__main__':
    continuous_results, binary_results, multi_results = data_analysis()
    # print('************************')
    # print(continuous_results[continuous_results["p_value"]<0.05]) #  "TotalWorkingYears","MonthlyIncome","Age","YearsAtCompany","YearsInCurrentRole","YearsWithCurrManager","整体满意度指数","在岗时间比例","YearsSinceLastPromotion","DistanceFromHome","入职收入增长率","换工作频率"
    # print(multi_results[multi_results["Chi2_p_value"]<0.05])  #       "StockOptionLevel","JobLevel","JobInvolvement","JobSatisfaction","EnvironmentSatisfaction","WorkLifeBalance",
    plot_feature_label_correlation(continuous_results, binary_results, multi_results)

# 'StockOptionLevel','JobLevel','MaritalStatus_Single','OverTime_Yes','isManager','Age','换工作频率','TotalWorkingYears','YearsWithCurrManager','入职收入增长率','YearsInCurrentRole','YearsAtCompany','整体满意度指数','MonthlyIncome',


# "MaritalStatus_Single","OverTime_Yes","isManager","StockOptionLevel","JobLevel","JobInvolvement","JobSatisfaction","EnvironmentSatisfaction","WorkLifeBalance","TotalWorkingYears","MonthlyIncome","Age","YearsAtCompany","YearsInCurrentRole","YearsWithCurrManager","整体满意度指数","在岗时间比例","YearsSinceLastPromotion","DistanceFromHome","入职收入增长率","换工作频率"
